/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once
#include <stack>
#include <mutex>

#include <baja/graphics/dx11.1/details/common.hpp>
#include <baja/graphics/dx11.1/graphics_fwd.hpp>
#include <baja/graphics/dx11.1/details/enum_mappings.hpp>
#include <baja/graphics/dx11.1/details/util.hpp>

#include <baja/win32/details/com_base.hpp>
#include <windows.ui.xaml.media.dxinterop.h>

namespace baja { namespace graphics { namespace dx11_1 {

namespace b = baja;
namespace d = baja::graphics::dx11_1::details;

BAJA_BEGIN_CONSTANTS(constants)
    BAJA_DECLARE_CONST(uint32, immediateModeVerticesPerBuffer, 1024)
BAJA_END_CONSTANTS

// constants
BAJA_BEGIN_GLOBALS(globals)
    BAJA_DECLARE_CONST_GLOBAL(wchar*, engineDisplayName)
    BAJA_DECLARE_CONST_GLOBAL(wchar*, warpEngineDisplayName)
    BAJA_DECLARE_CONST_GLOBAL(wchar*, featureLevelAddendum)
BAJA_END_GLOBALS(globals)

BAJA_DEFINE_CONST_GLOBAL(globals, wchar*, engineDisplayName, L"Direct3D 11")
BAJA_DEFINE_CONST_GLOBAL(globals, wchar*, warpEngineDisplayName, L"Direct3D 11 WARP")
BAJA_DEFINE_CONST_GLOBAL(globals, wchar*, featureLevelAddendum, L" - Feature Level ")

class engine : 
    public iengine, 
    public b::iunknown_provider,
    public std::enable_shared_from_this<iengine>
{
private:
    // used to work around make_shared only supporting up to 5 params
    struct create_params
    {
        create_params(
            IDXGISwapChain1* dxgiSwapChain_, 
            ID3D11Device1* d3d11Device_,
            ID3D11DeviceContext1* d3d11DeviceContext_,
            ID3D11DeviceContext1* d3d11DeviceContextDeferred_,
            D3D_FEATURE_LEVEL d3d11FeatureLevel_,
            std::wstring& engineName_,
            const display_params& displayParams_,
            uint32 msaaCount_,
            uint32 msaaQualityLevel_,
            bool isWarp_,
            ID2D1Device* d2d1Device_,
            ID2D1DeviceContext* d2d1DeviceContext_,
            ID2D1Factory1* d2d1Factory_
            ) :
                dxgiSwapChain(dxgiSwapChain_),
                d3d11Device(d3d11Device_),
                d3d11DeviceContext(d3d11DeviceContext_),
                d3d11DeviceContextDeferred(d3d11DeviceContextDeferred_),
                d3d11FeatureLevel(d3d11FeatureLevel_),
                engineName(engineName_),
                displayParams(displayParams_),
                msaaCount(msaaCount_),
                msaaQualityLevel(msaaQualityLevel_),
                isWarp(isWarp_),
                d2d1Device(d2d1Device_),
                d2d1DeviceContext(d2d1DeviceContext_),
                d2d1Factory(d2d1Factory_)
        {}

        IDXGISwapChain1* dxgiSwapChain; 
        ID3D11Device1* d3d11Device;
        ID3D11DeviceContext1* d3d11DeviceContext;
        ID3D11DeviceContext1* d3d11DeviceContextDeferred;
        D3D_FEATURE_LEVEL d3d11FeatureLevel;
        std::wstring engineName;
        display_params displayParams;
        uint32 msaaCount;
        uint32 msaaQualityLevel;
        bool isWarp;
        ID2D1Device* d2d1Device;
        ID2D1DeviceContext* d2d1DeviceContext;
        ID2D1Factory1* d2d1Factory;
    };

    // NOTE:
    // this is a hack we use to pass implementation-specific data through the abstraction layer by using a void* context
    // the caller code must pass an "equivalent" struct 
    // think about the better/cleaner way of doing this
    struct context_basic
    {
        shell_type shellType;
        IUnknown* window;

        context_basic(const void *context)
        {
            const context_basic* dest = static_cast<const context_basic*>(context);
            this->shellType = dest->shellType;
            this->window = dest->window;
        }
    };

    struct context_directui
    {
        shell_type shellType;
        WRL::ComPtr<ISwapChainBackgroundPanelNative> panelNative;

        context_directui(const void* context)
        {
            const context_directui* dest = static_cast<const context_directui*>(context);
            this->shellType = dest->shellType;
            this->panelNative = dest->panelNative;
        }
    };

    struct material_constants;
    struct light_constants;
    struct object_constants;
    struct misc_constants;
    struct immediate_mode;
    struct light_slot;
    struct stats;
    struct projection_matrix;
    struct shadow_data;

public:
    engine() {}

    engine(const engine::create_params& p);

    ~engine();
    
    static std::shared_ptr<engine> engine::create(
        const std::shared_ptr<iservice_map>& services,
        const bool useWARP,
        const display_params& dp,
        const void* context
        );

    display_params getDisplayParams();
    void changeDisplayParams(const display_params& displayParams);

    // resource creation
    bool getMultisampleLevels(const pixel_format format, uint32* msaaCount, uint32* msaaQualityLevel);

    std::shared_ptr<ieffect2> createEffect(const std::shared_ptr<effect2_description>& desc);

    std::shared_ptr<itexture> createTexture( 
        const uint32 width,
        const uint32 height,
        const uint32 mipLevels,
        const uint32 usage,
        const pixel_format format,
        const memory_pool pool,
        const uint32 msaaCount = 1,
        const uint32 msaaQuality = 0
        );

    concurrency::task<std::shared_ptr<g::itexture>>
    createTextureFromFileAsync(
        const wchar* filename,
        const uint32 usage,
        const pixel_format format,
        const memory_pool pool,
        const mipmap_load mipmapLoading
        );

    void
    createTextureFromFileAsync(
        const std::shared_ptr<baja::storage::istorage_file>& file,
        const uint32 usage,
        const pixel_format format,
        const memory_pool pool,
        const mipmap_load mipmapLoading,
        std::function<void(bool success, const std::shared_ptr<g::itexture>& texture)> createCompletedFunc
        );

    std::shared_ptr<itexture> createTextureFromFile(
        const wchar* filename,
        const uint32 usage,
        const pixel_format format,
        const memory_pool pool,
        const mipmap_load mipmapLoading
        );

    std::shared_ptr<itexture> engine::createTextureFromFileInMemory( 
        const memory* data,
        const uint32 dataSize,
        const uint32 usage,
        const pixel_format format,
        const memory_pool pool,
        const mipmap_load mipmapLoading
        );

    std::shared_ptr<isurface> createDepthStencilSurface( 
        const uint32 width,
        const uint32 height,
        const pixel_format format,
        const bool useFSAA
        );

    std::shared_ptr<isurface> createOffscreenSurface( 
        const uint32 width,
        const uint32 height,
        const pixel_format format
        );

    std::shared_ptr<isurface> createRenderTarget( 
        const uint32 width,
        const uint32 height,
        const pixel_format format,
        bool useFSAA
        );

    std::shared_ptr<ivertex_buffer> createVertexBuffer( 
        const uint32 totalBytes,
        const vertex_buffer_usage usage,
        const std::shared_ptr<iengine>& graphics
        );

    std::shared_ptr<iindex_buffer> createIndexBuffer(const uint32 totalIndices, const index_buffer_usage usage, const std::shared_ptr<g::iengine>& graphics);

    std::shared_ptr<ivertex_declaration> createVertexDeclaration();

    void setVertexDeclaration(const std::shared_ptr<ivertex_declaration>& vertexDecl, const uint32 context = 0);

    // render state management

    bool getIsHardwareAccelerated();
    void setAlphaBlend(const enabled_state state, const uint32 context = 0);
    void setAmbientLight(const color& color);
    void setBlendOp(const blend_op_state& state, const uint32 context = 0);
    void setCulling(const cull_state state, const uint32 context = 0);
    void setDstBlend(const blend_state state, const uint32 context = 0);
    void setFillMode(const fill_mode state, const uint32 context = 0);
    void setIndices(const std::shared_ptr<iindex_buffer>& indexBuffer, const uint32 context = 0);
    void setLight(const uint32 index, const light& light);
    void setLightEnable(const uint32 index, const enabled_state state);
    void setMaterial(const std::shared_ptr<imaterial>& material);
    void setSrcBlend(const blend_state state, const uint32 context = 0);
    void setStencil(const enabled_state state, const uint32 context = 0);
    void setStencilFail(const stencil_op state, const uint32 context = 0);
    void setStencilZFail(const stencil_op state, const uint32 context = 0);
    void setStencilPass(const stencil_op state, const uint32 context = 0);
    void setStencilCompare(const compare_state state, const uint32 context = 0);
    void setStencilRefValue(const uint32 refValue, const uint32 context = 0);
    void setStencilCompareMaskValue(const uint32 maskValue, const uint32 context = 0);
    void setStencilWriteMaskValue(const uint32 maskValue, const uint32 context = 0);

    void setStreamSource(
        const uint32 streamIndex,
        std::shared_ptr<ivertex_buffer>& streamData,
        const uint32 offsetInBytes,
        const uint32 vertexStride, 
        const uint32 context = 0
        );

    void setTexture(
        const uint32 textureIndex,
        const std::shared_ptr<itexture>& texture);

    void setTime(const float32 time);

    void getViewport(viewport* vp, const uint32 context = 0);
    void setViewport(const viewport& vp, const uint32 context = 0);

    void setZBias(const int32 bias, const uint32 context = 0);
    void setZBuffer(const enabled_state state, const uint32 context = 0);
    void setZCompare(const compare_state state, const uint32 context = 0);
    void setZWrite(const enabled_state state, const uint32 context = 0);

    void generateMips(const std::shared_ptr<itexture>& texture);

    std::shared_ptr<iasset_manager> getAssetManager();

    void clear(
        const clear_targets clearTargets,
        const packed_color packedColor,
        const float32 zValue,
        const uint32 stencilValue);

    std::shared_ptr<iscene_transformer> snapshotCurrentSceneTransforms();
    
    // render target / surface management

    std::shared_ptr<isurface> getRenderTarget(const uint32 context = 0);

    std::shared_ptr<isurface> getDepthStencilTarget(const uint32 context = 0);

    void setRenderTargets(
        const std::shared_ptr<isurface>& renderTarget,
        const std::shared_ptr<isurface>& depthStencilTarget,
        const uint32 context = 0
        );

    render_mode getRenderMode();
    void setRenderMode(const render_mode renderMode);

    void getBackBufferSize(uint32* width, uint32* height);

    // drawing

    void drawPrimitive(
        const primitive_type primitive,
        const uint32 startVertex,
        const uint32 primitiveCount,
        const uint32 context = 0
        );

    void drawIndexedPrimitive(
        const primitive_type primitive,
        const int32 baseVertexIndex,
        const uint32 minVertexIndex,
        const uint32 numVertices,
        const uint32 startIndex,
        const uint32 primitiveCount,
        const uint32 context = 0
        );

    // scene transformation management
    
    void pushWorldMatrix(const m::matrix44& world, const bool combineWithCurrentWorld);
    void popWorldMatrix();

    void pushCameraMatrix(const m::matrix44& camera);
    void replaceCurrentCameraMatrix(const m::matrix44& camera);
    void popCameraMatrix();

    void pushOrthographicProjection(
        const float32 left, 
        const float32 right, 
        const float32 top, 
        const float32 bottom, 
        const float32 nearZ, 
        const float32 farZ);

    void pushPerspectiveProjection(
        const float32 widthAtFocalPoint,
        const float32 heightAtFocalPoint, 
        const float32 nearZ,
        const float32 focalPointZ,
        const float32 farZ);

    void replaceCurrentProjectionWithOrthographicProjection(
        const float32 left, 
        const float32 right, 
        const float32 top, 
        const float32 bottom, 
        const float32 nearZ, 
        const float32 farZ);

    void replaceCurrentProjectionWithPerspectiveProjection(
        const float32 widthAtFocalPoint,
        const float32 heightAtFocalPoint, 
        const float32 nearZ,
        const float32 focalPointZ,
        const float32 farZ);

    void popProjectionMatrix();

    void pushUVMatrix(const m::matrix44& uvMatrix);
    void popUVMatrix();

    m::matrix44 getCoordinateSystemMatrix();
    void setCoordinateSystemMatrix(const m::matrix44& worldMatrix);

    m::matrix44 getWorldMatrix();
    m::matrix44 getWorldToProjectedMatrix();
    m::matrix44 getWorldToCameraMatrix();

    std::shared_ptr<ieffect2> getStandardEffect(const standard_effect standardEffect);

    void begin(const primitive_type primitive);
    void beginScreenSpace(const primitive_type primitive);

    void setColor(const packed_color packedColor);
    void setNormal(const m::vector3d& normal);
    void setTextureCoordinate(const texture_coordinate& texCoord);
    void setVertex(const m::point3d& position);
    void setScreenSpaceVertex(const m::point4d& position);

    void end();

    // scene management 

    void beginScene();
    void endScene();
    void showScene();
    void executeDeferred();

    // misc

    std::shared_ptr<ivar_group> getVarGroup();
    
    // d3d_wrapper methods
    size_t getInterfaceCount();
    WRL::ComPtr<IUnknown> getInterface(size_t index);

    // other public methods
    void setPixelShader(const WRL::ComPtr<ID3D11PixelShader>& pixelShader, const uint32 context = 0);
    void updateEffectTextures(const std::shared_ptr<ieffect2>& effect, const uint32 context = 0);
    
    void getCurrentBlendState(
        D3D11_BLEND_DESC* blendDesc, 
        float32* blendFactor, 
        uint32* sampleMask,
        const uint32 context = 0
        );

    void setBlendState(
        const D3D11_BLEND_DESC& blendDesc,
        const float32* blendFactor, 
        uint32 sampleMask,
        const uint32 context = 0
        );

    void getCurrentDepthStencilState(
        D3D11_DEPTH_STENCIL_DESC* depthStencilDesc, 
        uint32* stencilRef,
        const uint32 context = 0
        );

    void setDepthStencilState(
        const D3D11_DEPTH_STENCIL_DESC& depthStencilDesc,
        uint32 stencilRef,
        const uint32 context = 0
        );

    void getCurrentRasterizerState(
        D3D11_RASTERIZER_DESC* rasterizerDesc,
        const uint32 context = 0
        );

    void setRasterizerState(
        const D3D11_RASTERIZER_DESC& rasterizerDesc,
        const uint32 context = 0
        );

    WRL::ComPtr<ID3D11DeviceContext1> getContext(const uint32 context);

protected:

    void internalQueryInterface(const IID& iid, void** ppv);

private:

    void addStandardEffectDescription(
        const standard_effect e,
        const wchar* name,
        const wchar* pixelShaderFile,
        const wchar* vertexShaderFile
        );
    void initialize(const std::shared_ptr<iservice_map>& services);
    void setRenderTargetViews(const uint32 context = 0);
    void setD2DRenderTarget();
    void createNullTexture();
    void setCameraMatrix(const m::matrix44& camera);
    void setProjectionMatrix(const projection_matrix& projection);
    
    void updateMaterialConstantBuffer();
    void updateLightConstantBuffer();
    void updateObjectConstantBuffer();
    void updateMiscConstantBuffer(const uint32 context = 0);
    void updateShadowConstantBuffer();

    void initMembers(bool preserveState);
    void forceGraphicsIntoKnownState();
    void initializeDirectUI();

private:

    // constant buffer structs
    //
    // these structs use padding and different data types in places to adhere
    // to the shader constant's alignment issues (4 float vectors, 16 byte alignment, etc.)

    struct material_constants
    {
        material_constants()
        {
            ambient.set(1.0f,0.0f,0.0f,0.0f);
            diffuse.set(1.0f,1.0f,1.0f,1.0f);
            specular.set(0.0f, 0.0f, 0.0f, 0.0f);
            emissive.set(0.0f, 0.0f, 0.0f, 0.0f);
            specularPower = 1.0f;
            padding0 = 0.0f;
            padding1 = 0.0f;
            padding2 = 0.0f;
        }

        color ambient;
        color diffuse;
        color specular;
        color emissive;
        float32 specularPower;
        float32 padding0;
        float32 padding1;
        float32 padding2;
    };

    struct light_constants
    {
        light_constants()
        {
            ZeroMemory(this, sizeof(light_constants));
        }

        color ambient;
        color lightColor[g::constants::maxLights];
        float32 lightAttenuation[4*g::constants::maxLights];
        float32 lightDirection[4*g::constants::maxLights];
        float32 lightSpecularIntensity[4*g::constants::maxLights];
        uint32 isPointLight[4*g::constants::maxLights];
        uint32 activeLights;
        float32 padding0;
        float32 padding1;
        float32 padding2;
    };

    struct shadow_constants
    {
        shadow_constants()
        {
            ZeroMemory(this, sizeof(shadow_constants));
        }

        m::matrix44 shadowLocalToProjected4x4;
    };

   struct object_constants
    {
        object_constants()
        {
            ZeroMemory(this, sizeof(object_constants));
        }

        m::matrix44 localToWorld4x4;
        m::matrix44 localToProjected4x4;
        m::matrix44 worldToLocal4x4;
        m::matrix44 worldToView4x4;
        m::matrix44 uvTransform4x4;
        m::point3d eyePosition;
        float32 padding0;
    };

    struct misc_constants
    {
        misc_constants()
        {
            viewportWidth = 1.0f;
            viewportHeight = 1.0f;
            time = 0.0f;
            padding1 = 0.0f;
        }

        float32 viewportWidth;
        float32 viewportHeight;
        float32 time;
        float32 padding1;
    };

    // immediate mode struct
    struct immediate_mode
    {
        immediate_mode()
        {
            this->vertex.position = m::point3d();
            this->vertex.normal = m::vector3d(0.0f, 0.0f, 1.0f);
            this->vertex.color = packColor(0xff, 0xff, 0xff, 0xff);
            this->vertex.uv[0] = texture_coordinate();

            this->screenVertex.position = m::point4d();
            this->screenVertex.color = packColor(0xff, 0xff, 0xff, 0xff);
            this->screenVertex.uv[0] = texture_coordinate();

            this->isScreenSpace = false;
            this->isActive = false;
        }

        vertex_lit_colored_textured vertex;
        vertex_screen_colored_textured screenVertex;

        bool isScreenSpace;
        bool isActive;
    };

    // projection matrix struct
    struct projection_matrix 
    {
        projection_matrix()
        {
            this->minDepth = 1.0f;
            this->maxDepth = 1.0f;
        }

        m::matrix44 matrix;
        float32 minDepth;
        float32 maxDepth;
    };

    // shadow representation
    struct shadow_data 
    { 
        m::matrix44 matrix; 
        projection_matrix projection; 
    };

    // light slot struct
    struct light_slot
    {
        light_slot()
        {
            this->reset();
        }

        void reset()
        {
            ZeroMemory(&light, sizeof(light));
            enabled = false;
        }

        light light;
        bool enabled;
    };

    // stats
    struct stats
    {
        stats()
        {
            this->reset();
        }

        void reset()
        {
            ZeroMemory(this, sizeof(stats));
        }

        void maximize(const stats& updatedStats)
        {
            for (uint32 i = 0; i < static_cast<uint32>(primitive_type::triangleStrip) + 1; i++)
            {
                this->primitivesDrawn[i] = max(this->primitivesDrawn[i], updatedStats.primitivesDrawn[i]);
            }

            this->totalPrimitives = max(this->totalPrimitives, updatedStats.totalPrimitives);
            this->totalVertices = max(this->totalVertices, updatedStats.totalVertices);
        }

        uint32 primitivesDrawn[static_cast<uint32>(primitive_type::triangleStrip) + 1];
        uint32 totalPrimitives;
        uint32 totalVertices;
    };

    // template comparison functor for a block of memory
    template<typename T> struct memory_compare
    {
        // returns a < b
        bool operator()(const T& a, const T& b) const
        {
            return (std::memcmp(&a, &b, sizeof(T)) < 0);
        }
    };

    typedef std::map<D3D11_BLEND_DESC, WRL::ComPtr<ID3D11BlendState>, memory_compare<D3D11_BLEND_DESC>> 
        blend_state_map;
    typedef std::map<D3D11_DEPTH_STENCIL_DESC, WRL::ComPtr<ID3D11DepthStencilState>, memory_compare<D3D11_DEPTH_STENCIL_DESC>> 
        depth_stencil_state_map;
    typedef std::map<D3D11_RASTERIZER_DESC, WRL::ComPtr<ID3D11RasterizerState>, memory_compare<D3D11_RASTERIZER_DESC>> 
        rasterizer_state_map;

    blend_state_map _blendStateMap;
    depth_stencil_state_map _depthStencilStateMap;
    rasterizer_state_map _rasterizerStateMap;

    std::wstring _engineName;

    std::stack<m::matrix44> _worldMatrixStack;
    std::stack<m::matrix44> _worldInverseMatrixStack;
    std::stack<m::matrix44> _cameraMatrixStack;
    std::stack<projection_matrix> _projectionMatrixStack;
    std::stack<m::matrix44> _uvMatrixStack;

    // shadow stack
    std::stack<shadow_data> _shadowStack;

    m::matrix44 _coordinateSystemMatrix;

    display_params _displayParams;

    render_transforms _renderTransforms;
    color _ambientLight; 
    immediate_mode _immediateMode;
    stats _frameStats;
    stats _peakStats;

    D3D11_BLEND_DESC _defaultBlendDesc;
    D3D11_DEPTH_STENCIL_DESC _defaultDepthStencilDesc;
    D3D11_RASTERIZER_DESC _defaultRasterizerDesc;

    WRL::ComPtr<IDXGISwapChain1> _dxgiSwapChain;
    WRL::ComPtr<ID3D11Device1> _d3d11Device;
    WRL::ComPtr<ID3D11DeviceContext1> _d3d11DeviceContext;

    WRL::ComPtr<ID3D11DeviceContext1> _d3d11DeviceContextDeferred;

    // direct 2d
    WRL::ComPtr<ID2D1Factory1> _d2d1Factory;
    WRL::ComPtr<ID2D1Device> _d2d1Device;
    WRL::ComPtr<ID2D1DeviceContext> _d2d1DeviceContext;
    WRL::ComPtr<ID2D1DrawingStateBlock> _d2d1StateBlock;

    std::shared_ptr<idynamic_vertex_buffer> _worldIMBuffer;
    std::shared_ptr<idynamic_vertex_buffer> _screenIMBuffer;
    std::shared_ptr<ivar_group> _graphicsVarGroup;
    std::shared_ptr<iasset_manager> _assetManager;
    std::shared_ptr<effect2_description> _standardEffects[static_cast<uint32>(standard_effect::total)];
    WRL::ComPtr<ID3D11PixelShader> _overdrawShader;
    WRL::ComPtr<ID3D11PixelShader> _wireframeShader;
    std::shared_ptr<itexture> _nullTexture;
    std::shared_ptr<imaterial> _material;
    std::shared_ptr<itexture> _textures[g::constants::maxTextures];
    WRL::ComPtr<ID3D11Buffer> _materialBuffer;
    WRL::ComPtr<ID3D11Buffer> _lightBuffer;
    WRL::ComPtr<ID3D11Buffer> _objectBuffer;
    WRL::ComPtr<ID3D11Buffer> _miscBuffer;

    std::vector<light_slot> _lights;

    D3D_FEATURE_LEVEL _d3d11FeatureLevel;

    g::render_mode _currentRenderMode;
    g::render_mode _nextRenderMode;

    float32 _time;

    uint32 _msaaCount;
    uint32 _msaaQualityLevel;

    bool _isWarp;
    bool _resetPeakStats;
};

}}} // namespace baja.graphics.dx11_1
